import cv2 as cv
import numpy as np
import sys
from sklearn.cluster import KMeans

# load image
# path = r'D:\_const\large_data\pic\classmates_requests\Li_Yujie_2022_05_19Thu.png'
path = r'D:\_const\large_data\pic\classmates_requests\b1.png'
img = cv.imread(path, cv.IMREAD_GRAYSCALE)
ret, img = cv.threshold(img, 1, 255, cv.THRESH_BINARY)
img_ = img.copy()
# cv.imshow('img', img_)

# find seed point for floodFill
h, w = img.shape
mask = np.zeros((h, w), dtype=np.uint8)
mask[:, w//2] = 255
# cv.imshow('mask', mask)
res = cv.bitwise_and(img, img, mask= mask)
# cv.imshow('res', res)
res_middle = res[:, w//2]
idx = res_middle != 0
col_y = np.arange(0, h, dtype=int)
dots_y = col_y[idx]
print('dots_y', dots_y)
dots_y = dots_y.reshape(-1, 1)
kmeans = KMeans(n_clusters=3).fit(dots_y)
centers = kmeans.cluster_centers_
print('centers', centers)
centers = np.sort(centers, axis=0).ravel()
print('centers', centers)
point1 = (w // 2, int(np.round((centers[0] + centers[1]) / 2)))
point2 = (w // 2, int(np.round((centers[2] + centers[1]) / 2)))
print(point1, point2)
cv.circle(img, point1, 1, 255, 1)
cv.circle(img, point2, 1, 255, 1)
cv.imshow('points', img)
img = img_.copy()

# floodFill
# https://stackoverflow.com/questions/70322652/how-to-fill-color-correctly-in-closed-curve-using-opencv-python
img3c = np.stack((img, img, img), axis=2)
print('img3c', img3c.shape)
# cv.imshow('img3c', img3c)
mask = np.zeros((h + 2, w + 2), dtype=np.uint8)
cv.floodFill(img3c, mask, point1, (0, 255, 0), 1, 1)
cv.imshow('img filled', img3c)

# split and calculate
unit_x = w // 5
mask_tpl = np.zeros_like(img, dtype=np.uint8)
area_unit = w * h
print(f'Area unit = {area_unit}')
for i in range(5):
    number = i + 1
    print(f'----{number}----')
    xleft = i * unit_x
    mask = mask_tpl.copy()
    mask[:, xleft:xleft+unit_x] = 255
    res = cv.bitwise_and(img3c, img3c, mask=mask)
    idx = (res == (255, 255, 255))
    res[idx[:, :, 0]] = (0, 0, 0)
    ret, res = cv.threshold(res[:, :, 1], 127, 255, cv.THRESH_BINARY)
    cv.imshow(f'#{number}', res)
    # https://stackoverflow.com/questions/55467031/how-to-get-the-area-of-the-contours
    area = cv.countNonZero(res)
    print(f'Area = {area}, rate = {(area / area_unit):.4f}')

cv.waitKey()
cv.destroyAllWindows()
